//
//  main.cpp
//  opencv
//
//  Created by 徐欣 on 2016/11/17.
//  Copyright © 2016年 xx. All rights reserved.
//
#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/contrib/contrib.hpp>
#include <opencv2/stitching/stitcher.hpp>
#include <algorithm>
#include <random>
#include <time.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
#include "tinyxml2.h"

using namespace std;
using namespace cv;
using namespace tinyxml2;

class item {
public:
    int xmin;
    int ymax;
    Size size;
    string type;
};

vector<item> items;
vector<vector<item>> rows;

Mat slide_copy(Mat origin, Mat next, int height, int flag) {
    //滑动的距离
    int gap = 0;
    //两个图片的宽度之差
    int diff = abs(origin.cols - next.cols);
    if (diff != 0) {
        gap = rand() % diff;
    }
    if (origin.cols > next.cols) {
        vector<Mat> ch;
        split(next, ch);
        Mat ma(ch[3]);
        next.copyTo(Mat(origin, Rect(gap, height, next.cols, next.rows)),ma);
        
        if (flag!= -1 ) {
            for (int i = 0; i < rows[flag].size(); i++) {
                item *temp = new item;
                temp->xmin = rows[flag].at(i).xmin + gap;
                temp->size.width = rows[flag].at(i).size.width;
                temp->size.height = rows[flag].at(i).size.height;
                temp->ymax = height + next.rows;
                items.push_back(*temp);
            }
            //cout<<  "row"<< flag << ": " << rows[flag].size() <<endl;
        }
        
    } else {
        Rect rect(gap, 0, origin.cols, next.rows);
        Mat cut_next = Mat(next, rect);
        vector<Mat> ch;
        split(cut_next, ch);
        Mat ma(ch[3]);

        cut_next.copyTo(Mat(origin, Rect(0, height, cut_next.cols, cut_next.rows)),ma);
        
        if (flag!= -1 ) {
            for (int i = 0; i < rows[flag].size(); i++) {
                item a = rows[flag].at(i);
                if (a.xmin + 2 >= gap && (a.xmin + a.size.width -2 <= gap + origin.cols)) {
                    item *temp = new item;
                    temp->xmin = rows[flag].at(i).xmin - gap;
                    temp->size.width = rows[flag].at(i).size.width;
                    temp->size.height = rows[flag].at(i).size.height;
                    temp->ymax = height + next.rows;
                    items.push_back(*temp);
                }
            }
        }
        
    }
    return origin;
}

int main(int argc, char* argv[]) {
    Directory dir;
    
    string folder_name;
    string path;
    
    //cout << "input path:(eg:/Users/xuxin/Desktop/)" << endl;
    //cin >> path;
    //数据集放在执行文件的目录下
    path = "";
    cout << "input folder name:(eg:test)" << endl;
    cin >> folder_name;
    
    path += folder_name;
    cout<<path;
    
    //得到的货架的数量
    int count = 100;
    cout << "how many shelves do u want?(eg:100)" << endl;
    cin >> count;
    vector<Mat> result;
    
    //读取所有文件名称，存储->boxFileName
    vector<string> boxFileName;
    string exten = "*";
    bool addPath = false;
    vector<string> filenames = dir.GetListFiles(path, exten, addPath);
    cout << "the folder contains " << filenames.size() <<" files"<<endl;
    
    if (filenames.size() == 0) {
        return 0;
    }
    
    for (int t = 0; t < filenames.size(); t++) {
        string fullpath = path + "/" +filenames[t];
        boxFileName.push_back(fullpath);
    }
    
    //将所有的商品都resize为宽度接近120（改为随机）
    int resize_width = 120;
    
    for (int p = 0; p < count; p ++) {
        Size ns;
        //随机挑选n张商品素材，保存，计算宽度和高度
        vector<Mat> images;
        srand((unsigned)time(0)+ p * 400);
        //一排设定10到15件商品
        int n = rand()%5 + 10;
        Mat milk;
        for (int i = 0; i < n; i++) {
            do {
                int current = rand()%(boxFileName.size() - 1);
                milk = imread(boxFileName[current],-1);
            } while (!milk.data);
            
            Mat src;
            //resize_width - 8 < width < resize_width + 8
            int width = resize_width + (rand()%16 - 8);
            int height = ((double)width/milk.cols)*milk.rows;
            //预先处理
            resize(milk, milk, Size(width, height), (0, 0), (0, 0), cv::INTER_LINEAR);
            images.push_back(milk);
            //横向拼接
            ns.width += milk.cols;
            ns.height = milk.rows > ns.height? milk.rows: ns.height;
        }
        //加入随机空位 3个以内
        int vacancy = rand()%3;
        ns.width += resize_width*vacancy;
        // bgr
        Mat dst(ns, images[0].type(), Scalar(0, 0, 0));
        int width = 0;
        
        vector<item> row;
        for (int it = 0; it != images.size(); it ++) {
            item *temp = new item;
            temp->xmin = width;
            temp->size.width = images[it].cols;
            temp->size.height = images[it].rows;
            row.push_back(*temp);
            vector<Mat> channel;
            split(images[it], channel);
            Mat mask(channel[3]);
            images[it].copyTo(Mat(dst, Rect(width, dst.rows - images[it].rows, images[it].cols, images[it].rows)),mask);
            width += images[it].cols;
            if (vacancy > 0 && rand()%10 > 4) {
                vacancy --;
                width += resize_width;
            }
        }
        imwrite("/Users/xuxin/Desktop/o.jpg", dst);

        result.push_back(dst);
        rows.push_back(row);
    }
    
    
    
    //读入一张货架图片
    string shelf_jpg = "shelf/shelf.png";
    Mat shelf = imread(shelf_jpg, -1);
    if (!shelf.data) {
        cout<<"where is the shelf??? End."<<endl;
        return -1;
    }
    cvtColor(shelf,shelf,CV_RGB2BGRA);
    cout<<"shelf"<<shelf.channels()<<endl;

    mkdir((path + "/output_shelves").c_str(), 0775);
    mkdir((path + "/output_shelves/annotations").c_str(), 0775);
    //mkdir((path + "/output_trans").c_str(), 0775);
    //拼接货架
    for (int j = 0; j < count; j++) {
        stringstream stream;
        stream << j;
        string j_str = stream.str();
        string name = "/output_shelves/" + j_str;
        string xmlname = "/output_shelves/annotations/" + j_str;
        
        string shelfpath = path + name + "_shelves.jpg";
        string shelfxmlpath = path + xmlname + ".xml";
        
        //string rotate_name = "/output_trans/" + j_str+"rotate_shelves.jpg";
        //string rotate_path = path + rotate_name;
        
        //string perspective_name = "/trans/" + j_str+"Perspective_shelves.jpg";
        //string perspective_path = path + perspective_name;
        
        //纵向拼接 随机宽度
        Size ss;
        ss.width = (rand()%10 + 5) * resize_width;
        vector<int> choices;
        int choice_number = rand()%2 + 2;
        for (int z = 0; z < choice_number; z ++) {
            int temp = rand()%result.size();
            choices.push_back(temp);
        }
        
       
        ss.height = (shelf.rows+ 2) * choice_number;
        for (int zz = 0; zz < choice_number; zz++) {
            ss.height += result[choices[zz]].rows;
        }
        
        Mat shelves(ss, result[0].type(), Scalar::all(0));
        
        //处理shelf的width不够的问题
        int shelf_width = shelf.size().width;
        int shelf_number = 1;
        while (shelf_width < shelves.size().width) {
            shelf_width += shelf_width;
            shelf_number++;
        }
        Rect shelf_rect(0, 0, shelf_number*shelf.cols, shelf.rows);
        Mat shelf_all = Mat(shelf, shelf_rect);
        for (int x = 0; x < shelf_number; x++) {
            shelf.copyTo(Mat(shelf_all, Rect(x*shelf.cols,0, shelf.cols, shelf.rows)));
        }
        
        int current_height = 0;
        for(int zzz = 0; zzz < choice_number;zzz++) {
            shelves = slide_copy(shelves, result[choices[zzz]], current_height, choices[zzz]);
            current_height+=result[choices[zzz]].rows;
            shelves = slide_copy(shelves, shelf_all, current_height, -1);
            current_height+=shelf.rows+2;
        }
        
        imwrite(shelfpath, shelves);
    
        
        //生成xml文件
        XMLDocument* doc = new XMLDocument();
        doc->InsertEndChild(doc->NewDeclaration("xml version='1.0'"));
        XMLNode* anno = doc->InsertEndChild(doc->NewElement("annotation"));
        
        XMLElement* folder = doc->NewElement("folder");
        anno->InsertEndChild(folder);
        folder->SetText(folder_name.c_str());
        
        XMLElement* filename = doc->NewElement("filename");
        anno->InsertEndChild(filename);
        filename->SetText((j_str + "_shelves.jpg").c_str());
        
        XMLElement* size = doc->NewElement("size");
        XMLElement* w = doc->NewElement("width");
        size->InsertEndChild(w);
        w->SetText(shelves.cols);
        XMLElement* h = doc->NewElement("height");
        h->SetText(shelves.rows);
        size->InsertEndChild(h);
        
        anno->InsertEndChild(size);
        
        
        //cout << items.size() <<endl;
        for (int v = 0; v < items.size(); v++) {
            XMLElement* obj = doc->NewElement("object");
            
            XMLElement* obj_name = doc->NewElement("name");
            XMLElement* bndbox = doc->NewElement("bndbox");
            XMLElement* xmin = doc->NewElement("xmin");
            XMLElement* xmax = doc->NewElement("xmax");
            XMLElement* ymin = doc->NewElement("ymin");
            XMLElement* ymax = doc->NewElement("ymax");
            
            obj_name->SetText("box");
            
            xmin->SetText(items[v].xmin);
            xmax->SetText(items[v].xmin + items[v].size.width);
            ymin->SetText(items[v].ymax - items[v].size.height);
            //cout << "  "<<v <<" "<<items[v].ymax << "  "<< items[v].size.height << endl;
            ymax->SetText(items[v].ymax);
            
            bndbox->InsertEndChild(xmin);
            bndbox->InsertEndChild(xmax);
            bndbox->InsertEndChild(ymin);
            bndbox->InsertEndChild(ymax);
            
            obj->InsertEndChild(obj_name);
            obj->InsertEndChild(bndbox);
            
            anno->InsertEndChild(obj);
        }
        
        doc->SaveFile(shelfxmlpath.c_str());
        items.clear();
        
        
        //旋转中心为图像中心
        Point center = Point( shelves.cols/2, shelves.rows/2 );
        //计算二维旋转的仿射变换矩阵
        double angle = rand()%30 - 15;
        double scale = 1;
        Mat rotMat = getRotationMatrix2D(center, angle, scale);
        Mat shelves_rotate;
        //变换图像，并用黑色填充其余值
        warpAffine(shelves,shelves_rotate, rotMat, shelves.size());
        //imwrite(rotate_path, shelves_rotate);

        //rx 图像中心点 shelves.size
        //int rx = shelves.cols / 2;
        //int ry = shelves.rows / 2;
        
        /*
        x1= (x - rx0)*cos(a) - (y - ry0)*sin(a) + rx0 ;
        y1= (x - rx0)*sin(a) + (y - ry0)*cos(a) + ry0 ;
        */
        
        /*
        //透视变换 不使用
        int per_degree = 10;
        double sind = sin(per_degree);
        double cosd = cos(per_degree);
       
        Mat transform_matrix =(Mat_<double>(3,3) << 5, 0, 0, 0, cosd, 0, 0, -sind, 1);;
        warpPerspective(shelves, shelves_rotate, transform_matrix, shelves.size());
        imwrite(perspective_path, shelves_rotate);
        */
        
    }
    cout << "done!" <<endl;
    //waitKey(0);
    return 0;
}